import tkinter as tk
from tkinter import ttk, messagebox
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import sys
from collections import deque
import matplotlib

# 设置matplotlib使用支持中文的字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


class Edge:
    def __init__(self, to, cost, capacity, rev):
        self.to = to
        self.cost = cost
        self.capacity = capacity
        self.rev = rev


class MinCostFlowSolver:
    def __init__(self, n, edges):
        self.n = n
        self.edges = edges
        self.paths = [[], []]
        self.total_cost = 0

    def solve(self):
        # 初始化图
        max_nodes = self.n + 3
        e = [[] for _ in range(max_nodes)]

        def add_edge(u, v, cost, cap):
            e[u].append(Edge(v, cost, cap, len(e[v])))
            e[v].append(Edge(u, -cost, 0, len(e[u]) - 1))

        # 添加边
        for u, v, w in self.edges:
            add_edge(u, v, w, 1)
            add_edge(v, u, w, 1)

        # 添加超级源和汇
        s = self.n + 1
        t = self.n + 2
        add_edge(s, 1, 0, 2)  # 超级源到节点1，容量2
        add_edge(self.n, t, 0, 2)  # 节点n到超级汇，容量2

        INF = float('inf')
        path_count = 0

        while path_count < 2:
            # SPFA算法
            dist = [INF] * (self.n + 3)
            dist[s] = 0
            in_queue = [False] * (self.n + 3)
            q = deque([s])
            pre = [-1] * (self.n + 3)
            pre_edge = [0] * (self.n + 3)

            while q:
                u = q.popleft()
                in_queue[u] = False
                for i, edge in enumerate(e[u]):
                    if edge.capacity > 0 and dist[edge.to] > dist[u] + edge.cost:
                        dist[edge.to] = dist[u] + edge.cost
                        pre[edge.to] = u
                        pre_edge[edge.to] = i
                        if not in_queue[edge.to]:
                            in_queue[edge.to] = True
                            q.append(edge.to)

            if dist[t] == INF:
                break

            # 记录路径
            flow = INF
            v = t
            path = []
            while pre[v] != -1:
                u = pre[v]
                edge_index = pre_edge[v]
                flow = min(flow, e[u][edge_index].capacity)
                if 1 <= v <= self.n:  # 排除超级源和汇
                    path.append(v)
                v = u
            path.append(1)  # 添加起始节点
            path.reverse()

            # 去除连续重复节点
            unique_path = []
            last = -1
            for node in path:
                if node != last:
                    unique_path.append(node)
                    last = node

            # 修改点1：区分去程和回程路径
            if path_count == 0:  # 去程路径 (1→n)
                self.paths[0] = unique_path
            else:  # 回程路径 (n→1)
                self.paths[1] = unique_path[::-1]  # 关键修改：反转路径方向

            path_count += 1

            # 更新残余网络
            v = t
            while pre[v] != -1:
                u = pre[v]
                edge_index = pre_edge[v]
                e[u][edge_index].capacity -= flow
                e[v][e[u][edge_index].rev].capacity += flow
                v = u

            self.total_cost += dist[t] * flow

        return self.paths[0], self.paths[1], self.total_cost


class GraphVisualizer:
    def __init__(self, master, n, edges, path1, path2, total_cost):
        self.master = master
        self.n = n
        self.edges = edges
        self.path1 = path1
        self.path2 = path2
        self.total_cost = total_cost
        self.current_step = 0
        self.animation_speed = 500
        self.animation_running = False
        self.setup_ui()
        self.draw_graph()

    def setup_ui(self):
        self.master.title("无人机巡检路径规划可视化")
        # 主框架
        self.main_frame = ttk.Frame(self.master)
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        # 图形显示区域
        self.fig = plt.figure(figsize=(8, 6), dpi=100)
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.main_frame)
        self.canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=True)
        # 控制面板
        control_frame = ttk.Frame(self.main_frame)
        control_frame.pack(side=tk.BOTTOM, fill=tk.X)
        self.start_button = ttk.Button(control_frame, text="开始动画", command=self.start_animation)
        self.start_button.pack(side=tk.LEFT, padx=5, pady=5)
        self.stop_button = ttk.Button(control_frame, text="停止", command=self.stop_animation)
        self.stop_button.pack(side=tk.LEFT, padx=5, pady=5)
        self.step_button = ttk.Button(control_frame, text="单步", command=self.step_animation)
        self.step_button.pack(side=tk.LEFT, padx=5, pady=5)
        self.result_label = ttk.Label(control_frame, text=f"总路径长度: {self.total_cost}")
        self.result_label.pack(side=tk.RIGHT, padx=5, pady=5)

    def draw_graph(self):
        self.fig.clf()
        self.ax = self.fig.add_subplot(111)
        # 创建图
        self.G = nx.Graph()
        # 添加节点
        for i in range(1, self.n + 1):
            self.G.add_node(i)
        # 添加边
        for u, v, w in self.edges:
            self.G.add_edge(u, v, weight=w)

        # 计算节点位置
        if self.n <= 10:
            self.pos = nx.circular_layout(self.G)
        else:
            self.pos = nx.spring_layout(self.G, k=3 / (self.n ** 0.5), iterations=50)

        # 准备节点颜色列表
        node_colors = []
        for node in self.G.nodes():
            if node == 1:  # 起点1用绿色
                node_colors.append('green')
            elif node == self.n:  # 终点N用红色
                node_colors.append('red')
            else:  # 其他节点用浅蓝色
                node_colors.append('lightblue')

        # 绘制所有节点和边
        nx.draw_networkx_nodes(self.G, self.pos, node_size=500, node_color=node_colors, ax=self.ax)
        nx.draw_networkx_labels(self.G, self.pos, font_family='SimHei', ax=self.ax)
        nx.draw_networkx_edges(self.G, self.pos, edge_color='gray', width=1, ax=self.ax)

        # 绘制边权值
        edge_labels = {(u, v): w for u, v, w in self.edges}
        nx.draw_networkx_edge_labels(self.G, self.pos, edge_labels=edge_labels,
                                     font_family='SimHei', ax=self.ax)

        # 添加图例说明
        self.ax.set_title("无人机巡检网络图\n(绿色:起点1, 红色:终点N)", fontproperties='SimHei')
        self.ax.axis('off')
        self.canvas.draw()

    def start_animation(self):
        if not self.animation_running:
            self.animation_running = True
            self.animate()

    def stop_animation(self):
        self.animation_running = False

    def step_animation(self):
        self.stop_animation()
        self.update_animation()

    def animate(self):
        if self.animation_running and self.current_step < len(self.path1) + len(self.path2) + 2:
            self.update_animation()
            self.master.after(self.animation_speed, self.animate)
        else:
            self.animation_running = False

    def update_animation(self):
        self.fig.clf()
        self.ax = self.fig.add_subplot(111)

        # 准备节点颜色列表
        node_colors = []
        for node in self.G.nodes():
            if node == 1:  # 起点1用绿色
                node_colors.append('green')
            elif node == self.n:  # 终点N用红色
                node_colors.append('red')
            else:  # 其他节点用浅蓝色
                node_colors.append('lightblue')

        # 绘制基础图
        nx.draw_networkx_nodes(self.G, self.pos, node_size=500, node_color=node_colors, ax=self.ax)
        nx.draw_networkx_labels(self.G, self.pos, font_family='SimHei', ax=self.ax)
        nx.draw_networkx_edges(self.G, self.pos, edge_color='gray', width=1, ax=self.ax)

        # 绘制边权值
        edge_labels = {(u, v): w for u, v, w in self.edges}
        nx.draw_networkx_edge_labels(self.G, self.pos, edge_labels=edge_labels,
                                     font_family='SimHei', ax=self.ax)

        # 高亮已完成的路径
        if self.current_step > 0:
            # 绘制第一段路径 (去程 1→n)
            path1_edges = [(self.path1[i], self.path1[i + 1])
                           for i in range(min(self.current_step, len(self.path1) - 1))]
            nx.draw_networkx_edges(
                self.G, self.pos, edgelist=path1_edges,
                edge_color='red', width=3, ax=self.ax,
                arrows=True, arrowstyle='->'
            )

            # 如果第一段路径已完成，绘制第二段路径 (回程 n→1)
            if self.current_step > len(self.path1):
                # 关键修改：确保箭头方向正确
                path2_steps = min(self.current_step - len(self.path1), len(self.path2))
                path2_edges = []
                for i in range(path2_steps - 1):
                    # 注意这里使用path2[i]和path2[i+1]来保持方向一致
                    path2_edges.append((self.path2[i], self.path2[i + 1]))

                nx.draw_networkx_edges(
                    self.G, self.pos, edgelist=path2_edges,
                    edge_color='blue', width=3, ax=self.ax,
                    arrows=True, arrowstyle='->'
                )

        title = "无人机巡检网络图 - 路径规划中\n(绿色:起点1, 红色:终点N)" if self.current_step <= len(self.path1) + len(
            self.path2) else "无人机巡检网络图 - 规划完成\n(绿色:起点1, 红色:终点N)"
        self.ax.set_title(title, fontproperties='SimHei')
        self.ax.axis('off')
        self.canvas.draw()
        self.current_step += 1


class Application:
    def __init__(self, master):
        self.master = master
        self.master.title("无人机巡检最短往返路径规划系统")
        self.setup_ui()
        self.test_cases = []
        self.current_test_case = 0
        self.test_case_count = 0
        self.added_test_cases = 0

    def setup_ui(self):
        # 主框架
        self.main_frame = ttk.Frame(self.master)
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        # 输入区域
        input_frame = ttk.LabelFrame(self.main_frame, text="输入")
        input_frame.pack(fill=tk.X, padx=5, pady=5)
        ttk.Label(input_frame, text="测试用例数量 T:").grid(row=0, column=0, padx=5, pady=5)
        self.t_entry = ttk.Entry(input_frame)
        self.t_entry.grid(row=0, column=1, padx=5, pady=5)
        ttk.Button(input_frame, text="确认", command=self.set_test_case_count).grid(row=0, column=2, padx=5, pady=5)
        # 测试用例输入区域
        self.test_case_frame = ttk.LabelFrame(self.main_frame, text="测试用例输入")
        self.test_case_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        self.node_edge_label = ttk.Label(self.test_case_frame, text="节点数 N 和边数 M:(N M)")
        self.node_edge_label.pack(pady=5)
        self.node_edge_entry = ttk.Entry(self.test_case_frame)
        self.node_edge_entry.pack(pady=5)
        self.edges_label = ttk.Label(self.test_case_frame, text="边信息 (每行格式: 起点 终点 权值):")
        self.edges_label.pack(pady=5)
        self.edges_text = tk.Text(self.test_case_frame, height=10, width=50)
        self.edges_text.pack(pady=5, fill=tk.BOTH, expand=True)
        self.test_case_controls = ttk.Frame(self.test_case_frame)
        self.test_case_controls.pack(pady=5)
        ttk.Button(self.test_case_controls, text="添加测试用例", command=self.add_test_case).pack(side=tk.LEFT, padx=5)
        ttk.Button(self.test_case_controls, text="清空", command=self.clear_test_case).pack(side=tk.LEFT, padx=5)
        # 结果区域
        result_frame = ttk.LabelFrame(self.main_frame, text="结果")
        result_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        self.result_text = tk.Text(result_frame, height=10, width=50)
        self.result_text.pack(pady=5, fill=tk.BOTH, expand=True)
        # 控制按钮
        control_frame = ttk.Frame(self.main_frame)
        control_frame.pack(fill=tk.X, padx=5, pady=5)
        ttk.Button(control_frame, text="运行所有测试", command=self.run_all_tests).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="可视化当前测试", command=self.visualize_current_test).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="退出", command=self.master.quit).pack(side=tk.RIGHT, padx=5)

    def set_test_case_count(self):
        try:
            t = int(self.t_entry.get())
            if t <= 0:
                raise ValueError
            self.test_cases = []
            self.current_test_case = 0
            self.test_case_count = t
            self.added_test_cases = 0
            messagebox.showinfo("成功", f"已设置测试用例数量为 {t}")
        except ValueError:
            messagebox.showerror("错误", "请输入有效的测试用例数量(正整数)")

    def add_test_case(self):
        if self.test_case_count == 0:
            messagebox.showwarning("警告", "请先设置测试用例数量")
            return

        if self.added_test_cases >= self.test_case_count:
            messagebox.showwarning("警告", f"已达到设定的测试用例数量({self.test_case_count})")
            return

        try:
            n, m = map(int, self.node_edge_entry.get().split())
            if n <= 0 or m <= 0:
                raise ValueError("节点和边数量必须为正整数")

            edges_text = self.edges_text.get("1.0", tk.END).strip().split('\n')
            edges = []
            for line in edges_text:
                if line.strip():
                    u, v, w = map(int, line.split())
                    if u < 1 or u > n or v < 1 or v > n:
                        raise ValueError(f"节点编号必须在1到{n}之间")
                    edges.append((u, v, w))
            if len(edges) != m:
                raise ValueError(f"输入的边数量({len(edges)})与声明的边数量({m})不符")

            self.test_cases.append((n, m, edges))
            self.added_test_cases += 1
            self.clear_test_case()
            remaining = self.test_case_count - self.added_test_cases
            messagebox.showinfo("成功",
                                f"已添加测试用例 {self.added_test_cases}/{self.test_case_count}\n剩余可添加: {remaining}")
        except Exception as e:
            messagebox.showerror("错误", f"输入格式错误: {str(e)}")

    def clear_test_case(self):
        self.node_edge_entry.delete(0, tk.END)
        self.edges_text.delete("1.0", tk.END)

    def run_all_tests(self):
        if not self.test_cases:
            messagebox.showwarning("警告", "没有可运行的测试用例")
            return
        self.result_text.delete("1.0", tk.END)
        for idx, (n, m, edges) in enumerate(self.test_cases, 1):
            self.result_text.insert(tk.END, f"测试用例 {idx}:\n")
            self.result_text.insert(tk.END, f"节点数: {n}, 边数: {m}\n")
            try:
                solver = MinCostFlowSolver(n, edges)
                path1, path2, total_cost = solver.solve()

                if not path1 or not path2:
                    self.result_text.insert(tk.END, "无法找到有效的往返路径\n\n")
                    continue

                self.result_text.insert(tk.END, f"最短路径总长度: {total_cost}\n")
                self.result_text.insert(tk.END, f"去程路径 (1 -> {n}): {' -> '.join(map(str, path1))}\n")
                self.result_text.insert(tk.END, f"回程路径 ({n} -> 1): {' -> '.join(map(str, path2))}\n\n")
            except Exception as e:
                self.result_text.insert(tk.END, f"计算错误: {str(e)}\n\n")
        self.result_text.see(tk.END)

    def visualize_current_test(self):
        if not self.test_cases:
            messagebox.showwarning("警告", "没有可运行的测试用例")
            return
        if self.current_test_case >= len(self.test_cases):
            self.current_test_case = 0
        n, m, edges = self.test_cases[self.current_test_case]
        self.current_test_case += 1
        if n > 20:
            messagebox.showinfo("信息", "节点数超过20，不进行可视化，请查看文本结果")
            return
        try:
            solver = MinCostFlowSolver(n, edges)
            path1, path2, total_cost = solver.solve()
            if not path1 or not path2:
                messagebox.showerror("错误", "无法找到有效的往返路径")
                return
            # 在新窗口中显示可视化
            vis_window = tk.Toplevel(self.master)
            vis_window.title(f"测试用例 {self.current_test_case} 可视化")
            GraphVisualizer(vis_window, n, edges, path1, path2, total_cost)
        except Exception as e:
            messagebox.showerror("错误", f"可视化过程中出错: {str(e)}")


def main():
    root = tk.Tk()
    app = Application(root)
    root.mainloop()


if __name__ == "__main__":
    main()